#!/usr/bin/env python
# -*- coding: utf-8 -*-  
# Copyright (c)  2011 flykeysky<at>gmail.com. All rights reserved.
# This program or module is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License as published
# by the Free Software Foundation, either version 2 of the License, or
# version 3 of the License, or (at your option) any later version. It is
# provided for educational purposes and is distributed in the hope that
# it will be useful, but WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See
# the GNU General Public License for more details.

from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
from future_builtins import *

import os, sys
from PyQt4.QtCore import (QChar, QFile, QDir, QString, QTimer, QVariant, Qt,
        SIGNAL, QFileInfo)

from PyQt4.QtGui import (QApplication, QDialog, QHBoxLayout, QLabel,
        QListWidget, QListWidgetItem, QMessageBox, QPushButton,
        QSplitter, QTableWidget, QTableWidgetItem, QTreeWidget,
        QTreeWidgetItem, QVBoxLayout, QWidget, QWizard, QFileDialog)

import flows
import configwizard
import sigutils
#from configwizardui import Ui_ConfigWizard

MAC = True
try:
    from PyQt4.QtGui import qt_mac_set_native_menubar
except ImportError:
    MAC = False


class MainForm(QDialog):

    def __init__(self, parent=None):
        super(MainForm, self).__init__(parent)

        treeLabel = QLabel("Input Flows Tree")
        self.treeWidget = QTreeWidget()
        #self.treeWidget 
        self.treeWidget.setSelectionMode(self.treeWidget.ExtendedSelection) 
        self.treeWidget.setDragDropMode(self.treeWidget.InternalMove) 
        self.treeWidget.setDragEnabled(True) 
        self.treeWidget.setDropIndicatorShown(True) 
        treeLabel.setBuddy(self.treeWidget)
        
        tableLabel = QLabel("Selected Flows Table")
        self.tableWidget = QTableWidget()
        self.tableWidget.setSelectionBehavior(QTableWidget.SelectRows)
        tableLabel.setBuddy(self.tableWidget)
        
        configurewizardButton = QPushButton("Run &Config")
        extractButton = QPushButton("E&xtract flows")
        generatorwizardButton = QPushButton("&Generator")
        
        addFlowButton = QPushButton("&Add Flow")
        removeFlowButton = QPushButton("&Remove Flow")
        refreshFlowButton = QPushButton("Re&fresh")
        quitButton = QPushButton("&Quit")
        
        if not MAC:
            addFlowButton.setFocusPolicy(Qt.NoFocus)
            removeFlowButton.setFocusPolicy(Qt.NoFocus)
            quitButton.setFocusPolicy(Qt.NoFocus)
            configurewizardButton.setFocusPolicy(Qt.NoFocus)
            generatorwizardButton.setFocusPolicy(Qt.NoFocus)

        splitter = QSplitter(Qt.Horizontal)
        #tree
        vbox = QVBoxLayout()
        vbox.addWidget(treeLabel)
        vbox.addWidget(self.treeWidget)
        widget = QWidget()
        widget.setLayout(vbox)
        splitter.addWidget(widget)
        #table
        vbox = QVBoxLayout()
        vbox.addWidget(tableLabel)
        vbox.addWidget(self.tableWidget)
        widget = QWidget()
        widget.setLayout(vbox)
        splitter.addWidget(widget)

        buttonLayout = QHBoxLayout()
        buttonLayout.addWidget(configurewizardButton)
        buttonLayout.addWidget(extractButton)
        buttonLayout.addStretch()
        buttonLayout.addWidget(addFlowButton)
        buttonLayout.addStretch()
        buttonLayout.addWidget(removeFlowButton)
        buttonLayout.addWidget(refreshFlowButton)
        buttonLayout.addStretch()
        buttonLayout.addWidget(generatorwizardButton)
        buttonLayout.addWidget(quitButton)
        layout = QVBoxLayout()
        layout.addWidget(splitter)
        layout.addLayout(buttonLayout)
        self.setLayout(layout)

        #self.connect(self.tableWidget,SIGNAL("itemChanged(QTableWidgetItem*)"),  self.tableItemChanged)
        self.connect(configurewizardButton, SIGNAL("clicked()"), self.configurewiz)
        self.connect(extractButton, SIGNAL("clicked()"), self.extractflow)
        self.connect(addFlowButton, SIGNAL("clicked()"), self.addFlow)
        self.connect(removeFlowButton, SIGNAL("clicked()"), self.removeFlow)
        self.connect(refreshFlowButton, SIGNAL("clicked()"), self.refreshtable)
        self.connect(generatorwizardButton, SIGNAL("clicked()"), self.generatePayloadFileListFile)
        self.connect(quitButton, SIGNAL("clicked()"), self.accept)

        self.flows = flows.FlowContainer(QString("flows.dat"))
        self.selectedflows = flows.FlowContainer(QString("selectedflows.dat"))
        self.setWindowTitle("Flows (dict)")
        
        self.myGlobalConfig = {}
        self.PayloadFullname_list = []
        self.payloadfullnamelist_file = '/tmp/payloadfullnamelist'
        self.Payload5tuple_list = []
        QTimer.singleShot(0, self.initialLoad)


    def initialLoad(self):
        if not QFile.exists(self.flows.filename):
            #for flow in flows.generateFakeFlows():
                #self.flows.addFlow(flow)
            self.flows.dirty = False
        else:
            try:
                self.flows.load()
            except IOError, e:
                QMessageBox.warning(self, "Flows - Error",
                        "Failed to load: {0}".format(e))
        if not QFile.exists(self.selectedflows.filename):
            #for flow in flows.generateFakeFlows():
            #    self.selectedflows.addFlow(flow)
            self.selectedflows.dirty = False
        else:
            try:
                self.selectedflows.load()
            except IOError, e:
                QMessageBox.warning(self, "Flows - Error",
                        "Failed to load: {0}".format(e))
        self.myGlobalConfig = sigutils.LoadYAMLStream(configwizard.CONFIGFileName)
        self.populateTable()
        self.tableWidget.sortItems(0)
        self.populateTree()
        
    def configurewiz(self):
        myapp = configwizard.StartConfigWizard(self)
        myapp.show()
        
    def reject(self):
        self.accept()


    def accept(self):
        if ((self.selectedflows.dirty or self.flows.dirty) and
            QMessageBox.question(self, "Flows - Save?",
                    "Save unsaved changes?",
                    QMessageBox.Yes | QMessageBox.No) == 
                    QMessageBox.Yes):
            try:
                self.flows.save()
                self.selectedflows.save()
            except IOError, e:
                QMessageBox.warning(self, "Flows - Error",
                        "Failed to save: {0}".format(e))
        QDialog.accept(self)    
        
    def extractflow(self):
        #load yaml file
        #self.myGlobalConfig = sigutils.LoadYAMLStream(configwizard.CONFIGFileName)
        print("extracting...")
        print("%s\t%s\t%s" % (self.myGlobalConfig['ARFFFileName'], \
                                 self.myGlobalConfig['InputPayloadPath_s'], \
                                 self.myGlobalConfig['OutputSigDir_s']))
        if (self.myGlobalConfig['ARFFFileName'] == '' or \
                self.myGlobalConfig['InputPayloadPath_s'] == '' or \
                self.myGlobalConfig['OutputSigDir_s'] == ''):
            self.configurewiz()
        self.FindnLoadPayloadFiles()
        self.extractFromARFF()
    def FindnLoadPayloadFiles(self):
        m_PayloadFilesPath = self.myGlobalConfig['InputPayloadPath_s']
        #m_payloadfullname_list = []
        for (thisDir, subsHere, filesHere) in os.walk(str(m_PayloadFilesPath)):
            for filename in filesHere:
                (shortname, extension) = os.path.splitext(filename)
                payloadfullname = os.path.join(thisDir, filename)
                if(os.path.isfile(payloadfullname) and (extension == ".fldat")):
                    #print(payloadfullname)
                    self.PayloadFullname_list.append(payloadfullname)
                    self.Payload5tuple_list.append(shortname)
        #self.PayloadFullname_list = m_payloadfullname_list
        #print (self.Payload5tuple_list)
        f_name = self.payloadfullnamelist_file
        f = open(f_name, 'w')
        for i in self.PayloadFullname_list:
            #print(i)
            f.write(i + "\n")
        #sigutils.SaveYAMLStream(f_name,self.PayloadFullname_list)
        f.close()
    def extractFromARFF(self):
        import re
        attribute_prog = re.compile("^@ATTRIBUTE\ ")
        data_prog = re.compile("^@DATA")
        pARFFFile = open(self.myGlobalConfig['ARFFFileName'], 'r')
        feature_data = []
        all_index = {}
        feature_count = 0
        selectedfeature_names   =   ['max_fpktl', 'min_bpktl', 'std_bpktl', 'min_biat', 'sflow_fbytes', 'fpsh_cnt' ]
        selectedfivetupe_names  =   ['srcip', 'srcport', 'proto', 'dstip', 'dstport']
        selected_app_names      =   ['open_class'] #, 'app_class', 'cata_class']
        foundData               =   False
        flows_found_count       =   0
        flows_all_count         =   0
        flows_missed_count      =   0
        for line in pARFFFile.readlines():
            attribute_re = attribute_prog.search(line)
            if attribute_re:
                o = line.split()
                # to build a dict for 5-tuples, features and application info.
                print (feature_count, o[1])
                appendqueue_cmd = "all_index[\'%s\']= %d" % (str(o[1].strip()), feature_count)
                appendqueue_code = compile(appendqueue_cmd, '', 'exec')
                exec(appendqueue_code)
                feature_count += 1
                continue
            if 0:#foundData == False: #for debug
                for i in selectedfeature_names:
                    print(i)
                    if all_index.has_key(i):
                        print(i, all_index[i])
            m = data_prog.match(line)
            if m:
                foundData = True
            #print ("...")
            if (foundData == True and (\
                                        not line.isspace()) \
                                        and (not re.match('^@', line))  \
                                        and (not re.match('^%', line))  \
                                        ):
                l = line.strip().split(',') 
                #print(l)
                feature_data.append(l)
                m_opendpi =''
                m_l4 =''
                m_processname ='' 
                m_feature =''
                m_payloadfilename =''
                m_5tuple = ''
                #print (l[ all_index['proto'] ])
                flows_all_count += 1
                if (l[ all_index['proto'] ] == '17' or l[ all_index['proto'] ] == '6'):
                    count = 0
                    for ii in selectedfivetupe_names:
                        if ii == 'proto':
                            if l[ all_index['proto'] ] == '17':
                                m_5tuple += 'UDP' + '-'
                                m_l4 = 'UDP'
                            else:#l[ all_index['proto'] ] == '6'
                                m_5tuple += 'TCP' + '-'
                                m_l4 = 'TCP'
                        elif (count < len(selectedfivetupe_names) - 1):
                            m_5tuple += l[all_index[ii]] + '-'
                        else:
                            m_5tuple += l[all_index[ii]]
                        count += 1
                    print ( m_5tuple, len(selectedfivetupe_names) )
                    #
                    if (m_5tuple in self.Payload5tuple_list):
                        flows_found_count += 1
                        print ("---------%d flow found-----------" % (flows_found_count))
                        m_payloadfilename = self.myGlobalConfig['InputPayloadPath_s'] +'/'+ m_5tuple + '.fldat'
                        print (m_payloadfilename)
                        #print('payloadfile is: %s/%s.fldat' \
                        #      % (self.myGlobalConfig['InputPayloadPath_s'] , m_5tuple) )
                        count = 0
                        for ii in selectedfeature_names:
                            #print(ii, len(selectedfeature_names))
                            if (count < len(selectedfeature_names) -1):
                                m_feature += l[all_index[ii]] + ','
                            else:
                                m_feature += l[all_index[ii]]
                            count += 1
                        print (m_feature, len(selectedfeature_names) ) #<-- this the features
                        # extract opendpi
                        m_opendpi = l[all_index['open_class']]
                        print (m_opendpi)
                        m_opendpi_splited = m_opendpi.strip().split('_')
                        #print (m_opendpi_splited)
                        m_processname = m_opendpi_splited[-1]
                        print (m_processname)
                        flow = flows.Flow(m_opendpi, m_l4, m_processname, m_5tuple , m_feature, m_payloadfilename)
                        self.flows.addFlow(flow)
                    else:
                        print('no payload found.')
                        flows_missed_count +=1
        print ("--------------------\n%s found flows\n%s missed flows\n%s flows in all" % \
               (flows_found_count,flows_missed_count,flows_all_count))            
        pARFFFile.close()     
        self.populateTree()
    def populateTable(self, selectedFlow=None):
        selected = None
        self.tableWidget.clear()
        self.tableWidget.setSortingEnabled(True)
        self.tableWidget.setRowCount(len(self.selectedflows))
        headers = ["OpenDPI", "L4", "ProcessName", "FiveTuples", "FEATURES", "PayloadFileName"]
        self.tableWidget.setColumnCount(len(headers))
        self.tableWidget.setHorizontalHeaderLabels(headers)
        for row, flow in enumerate(self.selectedflows):
            item = QTableWidgetItem(flow.opendpi)
            item.setData(Qt.UserRole, QVariant(long(id(flow))))
            if selectedFlow is not None and selectedFlow == id(flow):
                selected = item
            self.tableWidget.setItem(row, flows.OPENDPI, item)
            self.tableWidget.setItem(row, flows.L4,
                    QTableWidgetItem(flow.l4))
            self.tableWidget.setItem(row, flows.PROCESSNAME,
                    QTableWidgetItem(flow.processname))
            self.tableWidget.setItem(row, flows.FIVETUPLE,
                    QTableWidgetItem(flow.fivetuple))
            self.tableWidget.setItem(row, flows.FEATURES,
                    QTableWidgetItem(flow.features))
            self.tableWidget.setItem(row, flows.PAYLOADFILENAME,
                    QTableWidgetItem(flow.payloadfilename))
        self.tableWidget.setSortingEnabled(True)
        self.tableWidget.resizeColumnsToContents()
        if selected is not None:
            selected.setSelected(True)
            self.tableWidget.setCurrentItem(selected)


    def populateTree(self, selectedFlow=None):
        selected = None
        self.treeWidget.clear()
        self.treeWidget.setColumnCount(3)
        self.treeWidget.setHeaderLabels(["ProcessName/L4/OpenDPI", "FEATURES", "PayloadFileName"])
        self.treeWidget.setItemsExpandable(True)
        parentFromProcessName = {}
        parentFromProcessNameL4 = {}
        parentFromProcessNameL4OpenDPI = {}
        for flow in self.flows.inProcessL4Order():
            ancestor = parentFromProcessName.get(flow.processname)
            if ancestor is None:
                ancestor = QTreeWidgetItem(self.treeWidget, [flow.processname])
                parentFromProcessName[flow.processname] = ancestor
            
            processnameL4 = flow.processname + "/" + flow.l4
            parent = parentFromProcessNameL4.get(processnameL4)
            if parent is None:
                parent = QTreeWidgetItem(ancestor, [flow.l4])
                parentFromProcessNameL4[processnameL4] = parent
            
            processnameL4OpenDPI = processnameL4 + "/" + flow.opendpi
            subparent = parentFromProcessNameL4OpenDPI.get(processnameL4OpenDPI)
            if subparent is None:
                subparent = QTreeWidgetItem(parent, [flow.opendpi])
                parentFromProcessNameL4OpenDPI[processnameL4OpenDPI] = subparent
            
            item = QTreeWidgetItem(subparent, [flow.fivetuple, QString("%L1").arg(flow.features) , \
                                               QString("%L1").arg(flow.payloadfilename)])
            item.setTextAlignment(1, Qt.AlignRight | Qt.AlignVCenter)
            if selectedFlow is not None and selectedFlow == id(flow):
                selected = item
            self.treeWidget.expandItem(subparent)
            self.treeWidget.expandItem(parent)
            self.treeWidget.expandItem(ancestor)
        self.treeWidget.resizeColumnToContents(0)
        self.treeWidget.resizeColumnToContents(1)
        if selected is not None:
            selected.setSelected(True)
            self.treeWidget.setCurrentItem(selected)
    def generatePayloadFileListFile(self):
        #print("generating payloadfile:%s..."%( str( self.myGlobalConfig['OutputSigDir_s'] ) ))
        fd = QFileDialog(self)
        m_outputfilename = fd.getSaveFileName( self,"Output file to...", self.myGlobalConfig['OutputSigDir_s'], "" )
        m_outputfile = open(m_outputfilename,'w')
        for i in self.selectedflows:
            print(i.payloadfilename)
            m_outputfile.write("%s\n"%(i.payloadfilename) )
        m_outputfile.close()
        fi = QFileInfo(m_outputfilename)
        out_path = fi.dir().absolutePath()
        cmd_1 = ( "d2o -l %s -o %s/data -x %s/offsets" % (m_outputfilename, out_path, out_path) )
        print(cmd_1)
        os.system(cmd_1)
        datafilename = ("%s/data" % (out_path) )
        print(datafilename)
        if ( not os.path.isfile(datafilename) ):
            cmd_2 = ( "mksary -i %s.ary %s" % (datafilename,datafilename) )
        #cmd_3 = ( "" % () )
    def addFlow(self):
        #flow = flows.Flow(" Unknown", " Unknown", " Unknown")
        for i in self.treeWidget.selectedItems ():
            if (i.childCount() > 0):
                print ("TODO: skipping, instead of select all children")
            else:
                print ("%d\t%s\t%s%s" % (i.childCount(), i.text(0), i.text(1), i.text(2)))
                print (i.parent().text(0), i.parent().parent().text(0), i.parent().parent().parent().text(0))
                
                m_fivetuple = QString (i.text(0))
                m_feature = QString(i.text(1))
                m_payloadfilename = i.text(2)
                m_opendpi = QString(i.parent().text(0))
                m_l4 = QString(i.parent().parent().text(0))
                m_processname = QString(i.parent().parent().parent().text(0)) 
                
                
                flow = flows.Flow(m_opendpi, m_l4, m_processname, m_fivetuple , m_feature, m_payloadfilename)
                flow.fivetuple = QString(i.text(0))
                
                self.selectedflows.addFlow(flow)
        self.populateTable(id(flow))
    
    
    
    def refreshtable(self):
        self.populateTable()

    def currentTableFlow(self):
        item = self.tableWidget.item(self.tableWidget.currentRow(), 0)
        if item is None:
            return None
        return self.selectedflows.flow(item.data(Qt.UserRole).toLongLong()[0])


    def removeFlow1(self):
        flow = self.currentTableFlow()
        if flow is None:
            return
        self.selectedflows.removeFlow(flow)
        self.populateTable()
    def removeFlow2(self):
        flowitemlist = self.tableWidget.selectedItems()
        #flow = self.currentTableFlow()
        for item in flowitemlist:
            if item is None:
                return
            fl = self.selectedflows.flow( item.data(Qt.UserRole).toLongLong()[0] )
            self.selectedflows.removeFlow(fl)
    def removeFlow(self):
        self.removeFlow2()
        self.populateTable()
        
def main():
    app = QApplication(sys.argv)
    form = MainForm()
    form.show()
    app.exec_()

# Process command-line arguments.
if __name__ == '__main__':
    main()
    exit()

